import torch
from torch import nn
from torch.nn import functional as f
from d2l import torch as d2l

class Inception(nn.Module):
    def __init__(self, in_channels,c1,c2,c3,c4,**kwargs):
        super(Inception,self).__init__(**kwargs)
        self.path1_1 = nn.Conv2d(in_channels,c1,kernel_size=1)
        
        self.path2_1 = nn.Conv2d(in_channels,c2[0],kernel_size=1)
        self.path2_2 = nn.Conv2d(c2[0],c2[1],kernel_size=3,padding=1)

        self.path3_1 = nn.Conv2d(in_channels,c3[0],kernel_size=1)
        self.path3_2 = nn.Conv2d(c3[0],c3[1],kernel_size=5,padding=1)

        self.path4_1 = nn.MaxPool2d(kernel_size=3,stride=1,padding=1)
        self.path4_2 = nn.Conv2d(in_channels,c4,kernel_size=1)
    
    def forward(self,x):
        path1 = f.relu(self.path1_1(x))
        path2 = f.relu(self.path2_2(f.relu(self.path2_1(x))))
        path3 = f.relu(self.path3_2(f.relu(self.path3_1(x))))
        path4 = f.relu(self.path4_2(self.path4_1(x)))
        return torch.cat((path1,path2,path3,path4),dim=1)

b1 = nn.Sequential(
    nn.Conv2d(1,64,kernel_size=7,stride=2,padding=1),nn.ReLU(),
    nn.MaxPool2d(kernel_size=3,stride=2,padding=1)
)
b2 = nn.Sequential(
    nn.Conv2d(64,64,kernel_size=1),nn.ReLU(),
    nn.Conv2d(64,192,kernel_size=3,padding=1),
    nn.MaxPool2d(kernel_size=3,stride=2,padding=1)
)
b3 = nn.Sequential(
    Inception(192,64,(96,128),(16,32),32),
    Inception(256,128,(128,192),(32,96),64),
    nn.MaxPool2d(kernel_size=3,stride=2,padding=1)
)
b4 = nn.Sequential(
    Inception(480,192,(96,208),(16,48),64),
    Inception(512,160,(112,224),(24,64),64),
    Inception(512,128,(128,256),(24,64),64),
    Inception(512,112,(144,288),(32,64),64),
    Inception(528,256,(160,320),(32,128),128),
    nn.MaxPool2d(kernel_size=3,stride=2,padding=1)
)
b5 = nn.Sequential(
    Inception(832,256,(160,320),(32,128),128),
    Inception(832,384,(192,384),(48,128),128),
    nn.AdaptiveAvgPool2d((1,1)),
    nn.Flatten()
)
net = nn.Sequential(b1,b2,b3,b4,b5,nn.Linear(1024,10))
